package com.scheduler.utils.thread;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.RemovalListener;
import com.google.common.cache.RemovalNotification;
import com.google.common.collect.Queues;
import lombok.extern.slf4j.Slf4j;

import java.util.LinkedHashMap;
import java.util.Optional;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

/**
 * 可变线程池工厂，可以给线程池起名字，按名字获取使用（比如按任务区分），但要注意区分粒度不宜太过细致。
 * 注意：线程池使用完成后需要 shutdown，否则无法释放，一直占用资源
 * 和 MutableThreadPoolExecutorFactory 实现不同，MutableThreadPoolExecutorFactory 会自动释放线程，线程池不可持有，而 MutableThreadPoolExecutorFactory2 不会自动释放线程，线程池可以持有。
 */
@Slf4j
public final class MutableThreadPoolExecutorFactory {

    private static volatile boolean running = true;

    /**
     * 已创建的线程池，自动过期，但会校验是否可以过期，如果还有任务运行，则会重新放入缓存
     */
    private static final Cache<String, MutableThreadPoolExecutor> threadPoolMaps = CacheBuilder.newBuilder()
            .expireAfterAccess(5, TimeUnit.MINUTES)
            .removalListener(new RemovalListener<String, MutableThreadPoolExecutor>() {
                @Override
                public void onRemoval(RemovalNotification<String, MutableThreadPoolExecutor> notification) {
                    // 回调的时候，实际已经清理掉了
                    MutableThreadPoolExecutor mutableThreadPoolExecutor = notification.getValue();
                    if (mutableThreadPoolExecutor == null) {
                        log.warn("缓存对象已经被销毁，无法打印日志， cacheKey={}", notification.getKey());
                        return;
                    }
                    String name = mutableThreadPoolExecutor.getName();
                    @SuppressWarnings("UnnecessaryLocalVariable") ThreadPoolExecutor pool = mutableThreadPoolExecutor;
                    if (pool.isTerminated()) {
                        log.info("ThreadPool_isTerminated, name={}, poolSize={}, coreSize={}, maxSize={}, activeCount={}, taskCount={}, queueSize={}, createTime={}", name, pool.getPoolSize(), pool.getCorePoolSize(), pool.getMaximumPoolSize(), pool.getActiveCount(), pool.getTaskCount(), pool.getQueue().size(), mutableThreadPoolExecutor.getCreateTime());
                        return;
                    }
                    if (pool.isShutdown()) {
                        log.info("ThreadPool_isShutdown, name={}, poolSize={}, coreSize={}, maxSize={}, activeCount={}, taskCount={}, queueSize={}, createTime={}", name, pool.getPoolSize(), pool.getCorePoolSize(), pool.getMaximumPoolSize(), pool.getActiveCount(), pool.getTaskCount(), pool.getQueue().size(), mutableThreadPoolExecutor.getCreateTime());
                        // 再放进去，等到isTerminated
                        threadPoolMaps.put(name, mutableThreadPoolExecutor);
                        return;
                    }
                    // 这里不会自动停止线程池，需要业务使用完成后手动shutdown掉，如果缓存过期的时候没有 shutdown 掉，就再放回去
                    log.info("ThreadPool_isRunning, name={}, poolSize={}, coreSize={}, maxSize={}, activeCount={}, taskCount={}, queueSize={}, createTime={}", name, pool.getPoolSize(), pool.getCorePoolSize(), pool.getMaximumPoolSize(), pool.getActiveCount(), pool.getTaskCount(), pool.getQueue().size(), mutableThreadPoolExecutor.getCreateTime());
                    threadPoolMaps.put(name, mutableThreadPoolExecutor);
                }
            }).build();

    public static boolean isRunning() {
        return running;
    }

    public static MutableThreadPoolExecutor getFixedBlockingPool(String uniqueName, Supplier<Integer> maxPoolSizeSupplier, int queueSize) {
        try {
            return threadPoolMaps.get(uniqueName, () -> newBlockingWaitPolicyThreadPool(maxPoolSizeSupplier, uniqueName, maxPoolSizeSupplier.get(), queueSize));
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public static MutableThreadPoolExecutor getFixedCallerRunsPool(String uniqueName, Supplier<Integer> maxPoolSizeSupplier, int queueSize) {
        try {
            return threadPoolMaps.get(uniqueName, () -> newCallerRunsPolicyThreadPool(maxPoolSizeSupplier, uniqueName, maxPoolSizeSupplier.get(), queueSize));
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public static MutableThreadPoolExecutor getFixedAbortPool(String uniqueName, Supplier<Integer> maxPoolSizeSupplier, int queueSize) {
        try {
            return threadPoolMaps.get(uniqueName, () -> newAbortPolicyThreadPool(maxPoolSizeSupplier, uniqueName, maxPoolSizeSupplier.get(), queueSize));
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public static MutableThreadPoolExecutor getFixedDiscardPool(String uniqueName, Supplier<Integer> maxPoolSizeSupplier, int queueSize) {
        try {
            return threadPoolMaps.get(uniqueName, () -> newDiscardPolicyThreadPool(maxPoolSizeSupplier, uniqueName, maxPoolSizeSupplier.get(), queueSize));
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public static void shutdown(String uniqueName) {
        MutableThreadPoolExecutor mutableThreadPoolExecutor = threadPoolMaps.getIfPresent(uniqueName);
        Optional.ofNullable(mutableThreadPoolExecutor).ifPresent(ThreadPoolExecutor::shutdown);
    }

    /**
     * 线程池工具
     */
    private static MutableThreadPoolExecutor newBlockingWaitPolicyThreadPool(Supplier<Integer> maxPoolSizeSupplier, String uniqueName, int maxPoolSize, int queueSize) {
        return new MutableThreadPoolExecutor(
                maxPoolSizeSupplier,
                uniqueName,
                queueSize <= 0 ? 0 : maxPoolSize,
                maxPoolSize,
                5,
                TimeUnit.MINUTES,
                queueSize <= 0 ? Queues.newSynchronousQueue() : Queues.newArrayBlockingQueue(queueSize),
                new NamedThreadFactory("MTPEF-" + uniqueName),
                new BlockingWaitPolicy(uniqueName));
    }

    /**
     * 线程池工具
     */
    private static MutableThreadPoolExecutor newCallerRunsPolicyThreadPool(Supplier<Integer> maxPoolSizeSupplier, String uniqueName, int maxPoolSize, int queueSize) {
        return new MutableThreadPoolExecutor(
                maxPoolSizeSupplier,
                uniqueName,
                queueSize <= 0 ? 0 : maxPoolSize,
                maxPoolSize,
                5,
                TimeUnit.MINUTES,
                queueSize <= 0 ? Queues.newSynchronousQueue() : Queues.newArrayBlockingQueue(queueSize),
                new NamedThreadFactory("MTPEF-" + uniqueName),
                new CallerRunsPolicy(uniqueName));
    }

    /**
     * 周期调度特殊线程池工具，使用拒绝策略，使用同步队列，不把任务放在队列中
     */
    private static MutableThreadPoolExecutor newAbortPolicyThreadPool(Supplier<Integer> maxPoolSizeSupplier, String uniqueName, int maxPoolSize, int queueSize) {
        return new MutableThreadPoolExecutor(
                maxPoolSizeSupplier,
                uniqueName,
                queueSize <= 0 ? 0 : maxPoolSize,
                maxPoolSize,
                5,
                TimeUnit.MINUTES,
                queueSize <= 0 ? Queues.newSynchronousQueue() : Queues.newArrayBlockingQueue(queueSize),
                new NamedThreadFactory("MTPEF-" + uniqueName),
                new AbortPolicy(uniqueName));
    }

    /**
     * 周期调度特殊线程池工具，使用拒绝策略，使用同步队列，不把任务放在队列中
     */
    private static MutableThreadPoolExecutor newDiscardPolicyThreadPool(Supplier<Integer> maxPoolSizeSupplier, String uniqueName, int maxPoolSize, int queueSize) {
        return new MutableThreadPoolExecutor(
                maxPoolSizeSupplier,
                uniqueName,
                queueSize <= 0 ? 0 : maxPoolSize,
                maxPoolSize,
                5,
                TimeUnit.MINUTES,
                queueSize <= 0 ? Queues.newSynchronousQueue() : Queues.newArrayBlockingQueue(queueSize),
                new NamedThreadFactory("MTPEF-" + uniqueName),
                new DiscardPolicy(uniqueName));
    }

    /**
     * 阻塞等待策略，该策略理论上和{@link ThreadPoolExecutor.CallerRunsPolicy}一样，但在部分场景下有区别，比如限流。
     * 在使用令牌桶限流的情况下，如果用CallerRunsPolicy会导致提交任务的线程在最后执行，令牌桶的实例已经在内存中使用（Guava的实现是直接sleep了，所以无法唤醒，改动影响较大），
     * 会导致提交任务线程一直阻塞，会阻塞很长时间，因此使用本策略，可以让提交任务的线程阻塞在队列上，而不是限流上，确保当限流调整后可以快速响应。
     */
    private static class BlockingWaitPolicy implements RejectedExecutionHandler {

        private final String threadPoolName;

        private final AtomicInteger triggerTimes = new AtomicInteger();

        public BlockingWaitPolicy(String threadPoolName) {
            this.threadPoolName = threadPoolName;
        }

        /**
         * Put task r to queue, wait until put success.
         *
         * @param r the runnable task requested to be executed
         * @param e the executor attempting to execute this task
         */
        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            try {
                long startTime = System.currentTimeMillis();
                if (triggerTimes.incrementAndGet() % 1000 == 0) {
                    // 每触发1000次打印一个日志
                    log.warn("触发线程池拒绝策略-阻塞策略1，name={}, triggerTimes={}", threadPoolName, triggerTimes.get());
                    if(triggerTimes.get()>0){
                        triggerTimes.set(0);
                    }
                }

                int times = 0;
                long rt = 0;

                boolean offer = false;
                // 如果没发放进去也没停止，则继续重放
                while (!offer && running) {
                    offer = e.getQueue().offer(r, 20, TimeUnit.MILLISECONDS);
                    times++;
                    // 如果总 RT 超过100ms，打印日志
                    rt = System.currentTimeMillis() - startTime;
                    if (rt > 100) {
                        log.warn("触发线程池拒绝策略-阻塞策略2，name={}, success={}, offerTimes={}, rt={}", threadPoolName, offer, times, rt);
                    }
                }
            } catch (InterruptedException ex) {
                throw new RejectedExecutionException("Task " + r + " rejected from " + e);
            }
        }
    }

    private static class CallerRunsPolicy extends ThreadPoolExecutor.CallerRunsPolicy {
        private final String threadPoolName;

        private final AtomicInteger triggerTimes = new AtomicInteger();

        public CallerRunsPolicy(String threadPoolName) {
            this.threadPoolName = threadPoolName;
        }

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            log.info("触发线程池拒绝策略-CallerRun策略, name={}, triggerTimes={}", threadPoolName, triggerTimes.get());
            super.rejectedExecution(r, e);
        }
    }

    private static class DiscardPolicy extends ThreadPoolExecutor.DiscardPolicy {
        private final String threadPoolName;

        private final AtomicInteger triggerTimes = new AtomicInteger();

        public DiscardPolicy(String threadPoolName) {
            this.threadPoolName = threadPoolName;
        }

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            if (triggerTimes.incrementAndGet() % 1000 == 0) {
                // 每触发1000次打印一个日志
                log.info("触发线程池拒绝策略-丢弃策略, name={}, triggerTimes={}", threadPoolName, triggerTimes.get());
                if(triggerTimes.get()>0){
                    triggerTimes.set(0);
                }
            }
        }
    }

    private static class AbortPolicy extends ThreadPoolExecutor.AbortPolicy {
        private final String threadPoolName;

        private final AtomicInteger triggerTimes = new AtomicInteger();

        public AbortPolicy(String threadPoolName) {
            this.threadPoolName = threadPoolName;
        }

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            log.info("触发线程池拒绝策略-拒绝策略, name={}, triggerTimes={}", threadPoolName, triggerTimes.get());
            super.rejectedExecution(r, e);
        }
    }

    private static boolean isRunning = true;

    static {
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            isRunning = false;
            log.info("Server停机关闭线程池开始");
            running = false;
            ConcurrentMap<String, MutableThreadPoolExecutor> poolMap = threadPoolMaps.asMap();
            // 关闭所有线程池
            poolMap.forEach((k, v) -> Optional.ofNullable(v).ifPresent(ThreadPoolExecutor::shutdown));
            // 检查并等待是否已经完全关闭
            poolMap.forEach((k, v) -> Optional.ofNullable(v).ifPresent(item -> awaitThreadPoolStopped(item, k)));

            log.info("Server停机关闭线程池完成");
        }));
    }

    /**
     * 等待线程池关闭，需要先调用shutdown，否则会导致等待过久，Spring 的实现会一个一个 shutdown，等待过久
     */
    private static void awaitThreadPoolStopped(ExecutorService threadPoolExecutor, String name) {
        try {
            log.info("Server等待线程池关闭，最大等待2分钟，name={}", name);
            boolean awaitTermination = threadPoolExecutor.awaitTermination(120, TimeUnit.SECONDS);
            if (awaitTermination) {
                log.info("Server线程池关闭成功，name={}", name);
            } else {
                log.warn("Server线程池关闭失败，name={}", name);
            }
        } catch (InterruptedException e) {
            log.error("Server等待线程池停止异常", e);
        }
    }

    private static final String POOL_NAME = "Mutable2All";

    static {
        // 自身简单监控
        ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
        AtomicInteger atomicInteger = new AtomicInteger();
        scheduledExecutorService.scheduleAtFixedRate(() -> {
            if (!isRunning) {
                scheduledExecutorService.shutdown();
                return;
            }
            try {
                LinkedHashMap<String, MutableThreadPoolExecutor> map = new LinkedHashMap<>(threadPoolMaps.asMap());
                int poolSize = map.values().stream().mapToInt(ThreadPoolExecutor::getPoolSize).sum();
                int coreSize = map.values().stream().mapToInt(ThreadPoolExecutor::getCorePoolSize).sum();
                int maxSize = map.values().stream().mapToInt(ThreadPoolExecutor::getMaximumPoolSize).sum();
                int activeCount = map.values().stream().mapToInt(ThreadPoolExecutor::getActiveCount).sum();
                long taskCount = map.values().stream().mapToLong(ThreadPoolExecutor::getTaskCount).sum();
                int queueSize = map.values().stream().mapToInt(v -> v.getQueue().size()).sum();
                // +1
                atomicInteger.incrementAndGet();
                // 采样日志
                if (atomicInteger.get() % 300 == 0) {
                    log.info("TreadPool_All2, count={}, poolSize={}, coreSize={}, maxSize={}, activeCount={}, taskCount={}, queueSize={}", map.size(), poolSize, coreSize, maxSize, activeCount, taskCount, queueSize);
                }
//                // 这里的监控需要上报到 Prometheus ，监控所有线程池数量和队列数量（暂时只需要监控总数，子线程池数量先不管），@羽明
//                // 重点关注 count, poolSize, activeCount, queueSize 四个指标
//                ThreadPoolMetrics.poolCount.labels(POOL_NAME, POOL_NAME).set(map.size());
//                ThreadPoolMetrics.poolSize.labels(POOL_NAME, POOL_NAME).set(poolSize);
//                ThreadPoolMetrics.poolActiveSize.labels(POOL_NAME, POOL_NAME).set(activeCount);
//                ThreadPoolMetrics.poolQueueSize.labels(POOL_NAME, POOL_NAME).set(queueSize);

                if (atomicInteger.get() % 300 == 0) {
                    map.values().forEach(pool -> log.info("ThreadPool_Single2, name={}, poolSize={}, coreSize={}, maxSize={}, activeCount={}, taskCount={}, queueSize={}, createTime={}", pool.getName(), pool.getPoolSize(), pool.getCorePoolSize(), pool.getMaximumPoolSize(), pool.getActiveCount(), pool.getTaskCount(), pool.getQueue().size(), pool.getCreateTime()));
                    if(atomicInteger.get()>0){
                        atomicInteger.set(0);
                    }
                }

            } catch (Throwable t) {
                log.error("监控异常", t);
            }
            try {
                LinkedHashMap<String, MutableThreadPoolExecutor> map = new LinkedHashMap<>(threadPoolMaps.asMap());
                map.values().forEach(MutableThreadPoolExecutor::renewFromConfigIfNecessary);
            } catch (Throwable t) {
                log.error("动态刷新线程池大小异常", t);
            }
        }, 5, 5, TimeUnit.SECONDS);
    }
}
